package felix.rule.js.script;

import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import delight.nashornsandbox.NashornSandbox;
import delight.nashornsandbox.NashornSandboxes;
import felix.rule.js.util.CustomerThreadFactory;
import felix.rule.js.util.ListeningExecutor;
import jdk.nashorn.api.scripting.NashornScriptEngineFactory;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import javax.script.Invocable;
import javax.script.ScriptEngine;
import javax.script.ScriptException;
import java.util.UUID;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
@Service
@ConditionalOnProperty(prefix = "js", value = "enable", havingValue = "true", matchIfMissing = false)
public class NashornJsInvokeServiceImpl extends AbstractJsInvokeService {

    private NashornSandbox sandbox;
    private ScriptEngine engine;
    private ExecutorService jsExecutorThreadPool;
    private ScheduledExecutorService timeoutExecutorService;

    private final AtomicInteger jsPushedMsgs = new AtomicInteger(0);
    private final AtomicInteger jsSuccessMsgs = new AtomicInteger(0);
    private final AtomicInteger jsEvalMsgs = new AtomicInteger(0);
    private final AtomicInteger jsFailedMsgs = new AtomicInteger(0);
    private final AtomicInteger jsTimeoutMsgs = new AtomicInteger(0);
    private final FutureCallback<UUID> evalCallback = new JsStatCallback<>(jsEvalMsgs, jsTimeoutMsgs, jsFailedMsgs);
    private final FutureCallback<Object> invokeCallback = new JsStatCallback<>(jsSuccessMsgs, jsTimeoutMsgs, jsFailedMsgs);

    @Autowired
    @Getter
    private JsExecutorServiceImpl jsExecutor;

    @Value("${js.max_requests_timeout:0}")
    private long maxRequestsTimeout;

    @Value("${js.stats.enabled:false}")
    private boolean statsEnabled;

    @Value("${js.use_js_sandbox:false}")
    private boolean useJsSandbox;

    @Value("${js.executor_thread_pool_size:10}")
    private int jsExecutorThreadPoolSize;

    @Value("${js.max_cpu_time:1000}")
    private long maxCpuTime;

    @Value("${js.max_errors:-1}")
    private int maxErrors;

    @Value("${js.remove_blacklists_threshold:1}")
    private int removeBlacklistsThreshold;

    @Scheduled(fixedDelayString = "${js.stats.print_interval_ms:10000}")
    public void printStats() {
        if (statsEnabled) {
            int pushedMsgs = jsPushedMsgs.getAndSet(0);
            int successMsgs = jsSuccessMsgs.getAndSet(0);
            int evalMsgs = jsEvalMsgs.getAndSet(0);
            int failed = jsFailedMsgs.getAndSet(0);
            int timedOut = jsTimeoutMsgs.getAndSet(0);
            if (pushedMsgs > 0 || successMsgs > 0 || evalMsgs > 0 || failed > 0 || timedOut > 0) {
                log.info("JS Invoke Stats: pushed [{}] received [{}] eval [{}] success [{}] failed [{}] timedOut [{}]",
                        pushedMsgs, successMsgs + evalMsgs, evalMsgs, successMsgs, failed, timedOut);
            }
        }
    }

    @PostConstruct
    public void init() {
        if (maxRequestsTimeout > 0) {
            timeoutExecutorService = Executors.newSingleThreadScheduledExecutor(CustomerThreadFactory.forName("nashorn-js-timeout"));
        }
        if (useJsSandbox) {
            sandbox = NashornSandboxes.create();
            jsExecutorThreadPool = Executors.newWorkStealingPool(jsExecutorThreadPoolSize);
            sandbox.setExecutor(jsExecutorThreadPool);
            sandbox.setMaxCPUTime(maxCpuTime);
            sandbox.allowNoBraces(false);
            sandbox.allowLoadFunctions(true);
            sandbox.setMaxPreparedStatements(30);
        } else {
            NashornScriptEngineFactory factory = new NashornScriptEngineFactory();
            engine = factory.getScriptEngine("--no-java");
        }
    }

    @PreDestroy
    public void stop() {
        if (jsExecutorThreadPool != null) {
            jsExecutorThreadPool.shutdownNow();
        }
        if (timeoutExecutorService != null) {
            timeoutExecutorService.shutdownNow();
        }
    }

    @Override
    protected ListenableFuture<UUID> doEval(UUID scriptId, String functionName, String jsScript) {
        jsPushedMsgs.incrementAndGet();
        ListenableFuture<UUID> result = jsExecutor.executeAsync(() -> {
            try {
                if (useJsSandbox) {
                    sandbox.eval(jsScript);
                } else {
                    engine.eval(jsScript);
                }
                scriptIdToNameMap.put(scriptId, functionName);
                return scriptId;
            } catch (Exception e) {
                log.warn("Failed to compile JS script: {}", e.getMessage(), e);
                throw new ExecutionException(e);
            }
        });
        if (maxRequestsTimeout > 0) {
            result = Futures.withTimeout(result, maxRequestsTimeout, TimeUnit.MILLISECONDS, timeoutExecutorService);
        }
        Futures.addCallback(result, evalCallback, jsExecutor);
        return result;
    }

    @Override
    protected ListenableFuture<Object> doInvokeFunction(UUID scriptId, String functionName, Object[] args) {
        jsPushedMsgs.incrementAndGet();
        ListenableFuture<Object> result = jsExecutor.executeAsync(() -> {
            try {
                Object obj;
                if (useJsSandbox) {
                    obj = sandbox.getSandboxedInvocable().invokeFunction(functionName, args);
                } else {
                    obj = ((Invocable) engine).invokeFunction(functionName, args);
                }
                onScriptExecutionSuccess(scriptId);
                return obj;
            } catch (Exception e) {
                onScriptExecutionError(scriptId);
                throw new ExecutionException(e);
            }
        });

        if (maxRequestsTimeout > 0) {
            result = Futures.withTimeout(result, maxRequestsTimeout, TimeUnit.MILLISECONDS, timeoutExecutorService);
        }
        Futures.addCallback(result, invokeCallback, jsExecutor);
        return result;
    }

    protected void onScriptExecutionSuccess(UUID scriptId) {
        whiteListedFunctions.computeIfAbsent(scriptId, key -> new AtomicInteger(0)).incrementAndGet();
        if (isWhiteListed(scriptId)) {
            blackListedFunctions.remove(scriptId);
        }
    }

    @Override
    protected void doRelease(UUID scriptId, String functionName) throws ScriptException {
        if (useJsSandbox) {
            sandbox.eval(functionName + " = undefined;");
        } else {
            engine.eval(functionName + " = undefined;");
        }
    }

    @Override
    protected int getMaxErrors() {
        return maxErrors;
    }

    private boolean isWhiteListed(UUID scriptId) {
        if (whiteListedFunctions.containsKey(scriptId)) {
            AtomicInteger successCount = whiteListedFunctions.get(scriptId);
            return successCount.get() >= removeBlacklistsThreshold;
        } else {
            return false;
        }
    }

}
